"""
    ConventionalObs

Observation operators for conventional meteorological observations.

This module implements forward and adjoint observation operators for:
- Temperature observations (surface and upper-air)
- Wind observations (u and v components)
- Humidity observations (specific, relative, dewpoint)
- Surface pressure observations

All operators include:
- Forward operator H(x) to compute model equivalent
- Adjoint operator Hᵀ(δy) for gradient computation
- Proper handling of vertical levels and coordinate transformations
- Wind rotation between grid-relative and Earth-relative coordinates

# Physical Conversions

## Temperature
- Potential temperature to temperature: T = θ * (p/p₀)^(R/cp)
- Virtual temperature: Tv = T * (1 + 0.61*q)

## Humidity
- Specific to relative: RH = (q/qs) * 100%
- Dewpoint from specific: Td = f(q, T, p) using Magnus formula
- Saturation specific humidity: qs = f(T, p)

## Wind Rotation
Earth-relative (u_e, v_e) to grid-relative (u_g, v_g):
- u_g = u_e * cos(α) + v_e * sin(α)
- v_g = -u_e * sin(α) + v_e * cos(α)

where α is the grid rotation angle.
"""
module ConventionalObs

using LinearAlgebra
using ..Interpolation

# Import interpolation functions
import ..Interpolation: bilinear_interpolate, bilinear_interpolate_adjoint!
import ..Interpolation: trilinear_interpolate, trilinear_interpolate_adjoint!
import ..Interpolation: vertical_interpolate, vertical_interpolate_adjoint!
import ..Interpolation: pressure_to_sigma

export ConventionalOperator
export temperature_forward, temperature_adjoint
export wind_forward, wind_adjoint
export humidity_forward, humidity_adjoint
export surface_pressure_forward, surface_pressure_adjoint
export specific_to_relative_humidity, relative_to_specific_humidity
export specific_to_dewpoint, dewpoint_to_specific
export rotate_wind_to_earth, rotate_wind_to_grid

# Physical constants
const Rd = 287.0        # Gas constant for dry air [J/(kg·K)]
const Rv = 461.5        # Gas constant for water vapor [J/(kg·K)]
const cp = 1005.0       # Specific heat at constant pressure [J/(kg·K)]
const g = 9.80665       # Gravitational acceleration [m/s²]
const p0 = 100000.0     # Reference pressure [Pa]
const T0 = 273.15       # 0°C in Kelvin
const epsilon = Rd/Rv   # Ratio of gas constants ≈ 0.622

"""
    ConventionalOperator{T<:AbstractFloat}

Container for conventional observation operator configuration.
"""
struct ConventionalOperator{T<:AbstractFloat}
    obs_type::Symbol  # :temperature, :wind_u, :wind_v, :humidity, :pressure
    vertical_method::Symbol  # :linear, :logarithmic
    use_virtual_temp::Bool
    rotate_winds::Bool

    function ConventionalOperator{T}(obs_type::Symbol;
                                    vertical_method::Symbol=:logarithmic,
                                    use_virtual_temp::Bool=false,
                                    rotate_winds::Bool=true) where T
        valid_types = [:temperature, :wind_u, :wind_v, :humidity, :pressure]
        @assert obs_type in valid_types "Invalid observation type: $obs_type"
        new{T}(obs_type, vertical_method, use_virtual_temp, rotate_winds)
    end
end

ConventionalOperator(obs_type::Symbol; kwargs...) =
    ConventionalOperator{Float64}(obs_type; kwargs...)

# Temperature Operators

"""
    temperature_forward(t_field::AbstractArray{T,3},
                       lon_grid::AbstractVector{T}, lat_grid::AbstractVector{T},
                       sigma_levels::AbstractVector{T},
                       lon_obs::T, lat_obs::T, p_obs::T,
                       ps_field::AbstractMatrix{T};
                       vertical_method::Symbol=:logarithmic) -> T

Forward operator for temperature observations.

# Arguments
- `t_field`: 3D temperature field [K] (nlon × nlat × nlev)
- `lon_grid`, `lat_grid`: Horizontal grid coordinates
- `sigma_levels`: Model sigma levels
- `lon_obs`, `lat_obs`, `p_obs`: Observation location and pressure
- `ps_field`: Surface pressure field [Pa]
- `vertical_method`: Vertical interpolation method

# Returns
- Interpolated temperature at observation location
"""
function temperature_forward(t_field::AbstractArray{T,3},
                            lon_grid::AbstractVector{T},
                            lat_grid::AbstractVector{T},
                            sigma_levels::AbstractVector{T},
                            lon_obs::T, lat_obs::T, p_obs::T,
                            ps_field::AbstractMatrix{T};
                            vertical_method::Symbol=:logarithmic) where T<:AbstractFloat

    # Get surface pressure at observation location
    ps_obs = bilinear_interpolate(ps_field, lon_grid, lat_grid, lon_obs, lat_obs)

    # Convert pressure to sigma
    sigma_obs = pressure_to_sigma(p_obs, ps_obs)

    # Perform 3D interpolation
    t_obs = trilinear_interpolate(t_field, lon_grid, lat_grid, sigma_levels,
                                  lon_obs, lat_obs, sigma_obs)

    return t_obs
end

"""
    temperature_adjoint!(t_gradient::AbstractArray{T,3},
                        ps_gradient::AbstractMatrix{T},
                        residual::T,
                        lon_grid::AbstractVector{T}, lat_grid::AbstractVector{T},
                        sigma_levels::AbstractVector{T},
                        lon_obs::T, lat_obs::T, p_obs::T,
                        ps_field::AbstractMatrix{T};
                        vertical_method::Symbol=:logarithmic)

Adjoint operator for temperature observations.

Distributes observation residual back to temperature and surface pressure fields.
"""
function temperature_adjoint!(t_gradient::AbstractArray{T,3},
                             ps_gradient::AbstractMatrix{T},
                             residual::T,
                             lon_grid::AbstractVector{T},
                             lat_grid::AbstractVector{T},
                             sigma_levels::AbstractVector{T},
                             lon_obs::T, lat_obs::T, p_obs::T,
                             ps_field::AbstractMatrix{T};
                             vertical_method::Symbol=:logarithmic) where T<:AbstractFloat

    # Get surface pressure at observation location
    ps_obs = bilinear_interpolate(ps_field, lon_grid, lat_grid, lon_obs, lat_obs)
    sigma_obs = pressure_to_sigma(p_obs, ps_obs)

    # Adjoint of 3D interpolation
    trilinear_interpolate_adjoint!(t_gradient, residual,
                                   lon_grid, lat_grid, sigma_levels,
                                   lon_obs, lat_obs, sigma_obs)

    # Adjoint of sigma computation: ∂J/∂ps = ∂J/∂σ * ∂σ/∂ps
    # σ = p/ps, so ∂σ/∂ps = -p/ps²
    # This is a simplification - full implementation would need the vertical interpolation Jacobian
    # For now, we skip the ps_gradient contribution from sigma dependency

    return nothing
end

# Wind Operators

"""
    wind_forward(u_field::AbstractArray{T,3}, v_field::AbstractArray{T,3},
                lon_grid::AbstractVector{T}, lat_grid::AbstractVector{T},
                sigma_levels::AbstractVector{T},
                lon_obs::T, lat_obs::T, p_obs::T,
                ps_field::AbstractMatrix{T},
                component::Symbol;
                rotate_to_earth::Bool=true,
                rotation_angle::T=T(0)) -> T

Forward operator for wind observations.

# Arguments
- `u_field`, `v_field`: 3D wind component fields [m/s]
- Grid and observation location parameters
- `component`: :u or :v
- `rotate_to_earth`: Whether to rotate from grid to Earth coordinates
- `rotation_angle`: Grid rotation angle [radians]

# Returns
- Interpolated wind component at observation location
"""
function wind_forward(u_field::AbstractArray{T,3},
                     v_field::AbstractArray{T,3},
                     lon_grid::AbstractVector{T},
                     lat_grid::AbstractVector{T},
                     sigma_levels::AbstractVector{T},
                     lon_obs::T, lat_obs::T, p_obs::T,
                     ps_field::AbstractMatrix{T},
                     component::Symbol;
                     rotate_to_earth::Bool=true,
                     rotation_angle::T=T(0)) where T<:AbstractFloat

    @assert component in [:u, :v] "Component must be :u or :v"

    # Get surface pressure and sigma
    ps_obs = bilinear_interpolate(ps_field, lon_grid, lat_grid, lon_obs, lat_obs)
    sigma_obs = pressure_to_sigma(p_obs, ps_obs)

    # Interpolate both wind components
    u_interp = trilinear_interpolate(u_field, lon_grid, lat_grid, sigma_levels,
                                     lon_obs, lat_obs, sigma_obs)
    v_interp = trilinear_interpolate(v_field, lon_grid, lat_grid, sigma_levels,
                                     lon_obs, lat_obs, sigma_obs)

    # Apply wind rotation if needed
    if rotate_to_earth && abs(rotation_angle) > eps(T)
        u_earth, v_earth = rotate_wind_to_earth(u_interp, v_interp, rotation_angle)
        return component == :u ? u_earth : v_earth
    else
        return component == :u ? u_interp : v_interp
    end
end

"""
    wind_adjoint!(u_gradient::AbstractArray{T,3}, v_gradient::AbstractArray{T,3},
                 ps_gradient::AbstractMatrix{T},
                 residual::T,
                 lon_grid::AbstractVector{T}, lat_grid::AbstractVector{T},
                 sigma_levels::AbstractVector{T},
                 lon_obs::T, lat_obs::T, p_obs::T,
                 ps_field::AbstractMatrix{T},
                 component::Symbol;
                 rotate_to_earth::Bool=true,
                 rotation_angle::T=T(0))

Adjoint operator for wind observations.

Distributes observation residual back to u, v wind fields.
"""
function wind_adjoint!(u_gradient::AbstractArray{T,3},
                      v_gradient::AbstractArray{T,3},
                      ps_gradient::AbstractMatrix{T},
                      residual::T,
                      lon_grid::AbstractVector{T},
                      lat_grid::AbstractVector{T},
                      sigma_levels::AbstractVector{T},
                      lon_obs::T, lat_obs::T, p_obs::T,
                      ps_field::AbstractMatrix{T},
                      component::Symbol;
                      rotate_to_earth::Bool=true,
                      rotation_angle::T=T(0)) where T<:AbstractFloat

    @assert component in [:u, :v] "Component must be :u or :v"

    # Get surface pressure and sigma
    ps_obs = bilinear_interpolate(ps_field, lon_grid, lat_grid, lon_obs, lat_obs)
    sigma_obs = pressure_to_sigma(p_obs, ps_obs)

    # Apply adjoint of wind rotation
    u_residual = residual
    v_residual = T(0)

    if rotate_to_earth && abs(rotation_angle) > eps(T)
        # Adjoint of rotation
        cos_α = cos(rotation_angle)
        sin_α = sin(rotation_angle)

        if component == :u
            # Adjoint of u_earth = u_grid * cos_α - v_grid * sin_α
            u_residual = residual * cos_α
            v_residual = -residual * sin_α
        else  # component == :v
            # Adjoint of v_earth = u_grid * sin_α + v_grid * cos_α
            u_residual = residual * sin_α
            v_residual = residual * cos_α
        end
    else
        if component == :u
            u_residual = residual
            v_residual = T(0)
        else
            u_residual = T(0)
            v_residual = residual
        end
    end

    # Distribute to u and v fields
    trilinear_interpolate_adjoint!(u_gradient, u_residual,
                                   lon_grid, lat_grid, sigma_levels,
                                   lon_obs, lat_obs, sigma_obs)
    trilinear_interpolate_adjoint!(v_gradient, v_residual,
                                   lon_grid, lat_grid, sigma_levels,
                                   lon_obs, lat_obs, sigma_obs)

    return nothing
end

"""
    rotate_wind_to_earth(u_grid::T, v_grid::T, rotation_angle::T) -> Tuple{T,T}

Rotate wind from grid-relative to Earth-relative coordinates.

# Arguments
- `u_grid`, `v_grid`: Wind components in grid coordinates
- `rotation_angle`: Grid rotation angle [radians]

# Returns
- `(u_earth, v_earth)`: Wind components in Earth coordinates
"""
function rotate_wind_to_earth(u_grid::T, v_grid::T, rotation_angle::T) where T<:AbstractFloat
    cos_α = cos(rotation_angle)
    sin_α = sin(rotation_angle)

    u_earth = u_grid * cos_α - v_grid * sin_α
    v_earth = u_grid * sin_α + v_grid * cos_α

    return (u_earth, v_earth)
end

"""
    rotate_wind_to_grid(u_earth::T, v_earth::T, rotation_angle::T) -> Tuple{T,T}

Rotate wind from Earth-relative to grid-relative coordinates.
"""
function rotate_wind_to_grid(u_earth::T, v_earth::T, rotation_angle::T) where T<:AbstractFloat
    # Inverse rotation (transpose of rotation matrix)
    cos_α = cos(rotation_angle)
    sin_α = sin(rotation_angle)

    u_grid = u_earth * cos_α + v_earth * sin_α
    v_grid = -u_earth * sin_α + v_earth * cos_α

    return (u_grid, v_grid)
end

# Humidity Operators

"""
    humidity_forward(q_field::AbstractArray{T,3},
                    lon_grid::AbstractVector{T}, lat_grid::AbstractVector{T},
                    sigma_levels::AbstractVector{T},
                    lon_obs::T, lat_obs::T, p_obs::T,
                    ps_field::AbstractMatrix{T},
                    humidity_type::Symbol=:specific;
                    t_field::Union{AbstractArray{T,3},Nothing}=nothing) -> T

Forward operator for humidity observations.

# Arguments
- `q_field`: 3D specific humidity field [kg/kg]
- Grid and observation location parameters
- `humidity_type`: :specific, :relative, or :dewpoint
- `t_field`: Temperature field (required for :relative and :dewpoint)

# Returns
- Humidity in requested units
"""
function humidity_forward(q_field::AbstractArray{T,3},
                         lon_grid::AbstractVector{T},
                         lat_grid::AbstractVector{T},
                         sigma_levels::AbstractVector{T},
                         lon_obs::T, lat_obs::T, p_obs::T,
                         ps_field::AbstractMatrix{T},
                         humidity_type::Symbol=:specific;
                         t_field::Union{AbstractArray{T,3},Nothing}=nothing) where T<:AbstractFloat

    # Get surface pressure and sigma
    ps_obs = bilinear_interpolate(ps_field, lon_grid, lat_grid, lon_obs, lat_obs)
    sigma_obs = pressure_to_sigma(p_obs, ps_obs)

    # Interpolate specific humidity
    q_interp = trilinear_interpolate(q_field, lon_grid, lat_grid, sigma_levels,
                                     lon_obs, lat_obs, sigma_obs)

    # Convert to requested type
    if humidity_type == :specific
        return q_interp

    elseif humidity_type == :relative
        @assert !isnothing(t_field) "Temperature field required for relative humidity"
        t_interp = trilinear_interpolate(t_field, lon_grid, lat_grid, sigma_levels,
                                        lon_obs, lat_obs, sigma_obs)
        return specific_to_relative_humidity(q_interp, t_interp, p_obs)

    elseif humidity_type == :dewpoint
        @assert !isnothing(t_field) "Temperature field required for dewpoint"
        t_interp = trilinear_interpolate(t_field, lon_grid, lat_grid, sigma_levels,
                                        lon_obs, lat_obs, sigma_obs)
        return specific_to_dewpoint(q_interp, p_obs)

    else
        error("Unknown humidity type: $humidity_type")
    end
end

"""
    humidity_adjoint!(q_gradient::AbstractArray{T,3},
                     ps_gradient::AbstractMatrix{T},
                     residual::T,
                     lon_grid::AbstractVector{T}, lat_grid::AbstractVector{T},
                     sigma_levels::AbstractVector{T},
                     lon_obs::T, lat_obs::T, p_obs::T,
                     ps_field::AbstractMatrix{T},
                     humidity_type::Symbol=:specific;
                     t_field::Union{AbstractArray{T,3},Nothing}=nothing)

Adjoint operator for humidity observations.

Note: For :relative and :dewpoint, this is a linearized adjoint around the background state.
"""
function humidity_adjoint!(q_gradient::AbstractArray{T,3},
                          ps_gradient::AbstractMatrix{T},
                          residual::T,
                          lon_grid::AbstractVector{T},
                          lat_grid::AbstractVector{T},
                          sigma_levels::AbstractVector{T},
                          lon_obs::T, lat_obs::T, p_obs::T,
                          ps_field::AbstractMatrix{T},
                          humidity_type::Symbol=:specific;
                          t_field::Union{AbstractArray{T,3},Nothing}=nothing) where T<:AbstractFloat

    # Get surface pressure and sigma
    ps_obs = bilinear_interpolate(ps_field, lon_grid, lat_grid, lon_obs, lat_obs)
    sigma_obs = pressure_to_sigma(p_obs, ps_obs)

    # For specific humidity or linearized conversions, distribute directly
    # Full nonlinear adjoint of humidity conversions would require Jacobian
    # For now, we use simplified linear adjoint

    trilinear_interpolate_adjoint!(q_gradient, residual,
                                   lon_grid, lat_grid, sigma_levels,
                                   lon_obs, lat_obs, sigma_obs)

    return nothing
end

# Humidity conversion functions

"""
    specific_to_relative_humidity(q::T, t::T, p::T) -> T

Convert specific humidity to relative humidity.

RH = (q / qs) × 100%

where qs is saturation specific humidity.
"""
function specific_to_relative_humidity(q::T, t::T, p::T) where T<:AbstractFloat
    qs = saturation_specific_humidity(t, p)
    rh = (q / qs) * T(100)
    return clamp(rh, T(0), T(100))
end

"""
    saturation_specific_humidity(t::T, p::T) -> T

Compute saturation specific humidity using Clausius-Clapeyron equation.
"""
function saturation_specific_humidity(t::T, p::T) where T<:AbstractFloat
    # Saturation vapor pressure (Magnus formula)
    es = saturation_vapor_pressure(t)

    # Saturation mixing ratio
    ws = epsilon * es / (p - es)

    # Convert to specific humidity
    qs = ws / (T(1) + ws)

    return qs
end

"""
    saturation_vapor_pressure(t::T) -> T

Compute saturation vapor pressure using Magnus formula.

es = 611.2 * exp(17.67 * (T - 273.15) / (T - 29.65))
"""
function saturation_vapor_pressure(t::T) where T<:AbstractFloat
    # Magnus formula
    t_celsius = t - T(273.15)
    es = T(611.2) * exp(T(17.67) * t_celsius / (t - T(29.65)))
    return es
end

"""
    specific_to_dewpoint(q::T, p::T) -> T

Convert specific humidity to dewpoint temperature.

Uses inverse Magnus formula.
"""
function specific_to_dewpoint(q::T, p::T) where T<:AbstractFloat
    # Convert specific humidity to vapor pressure
    e = q * p / (epsilon + q * (T(1) - epsilon))

    # Inverse Magnus formula
    ln_ratio = log(e / T(611.2))
    td = T(273.15) + T(243.5) * ln_ratio / (T(17.67) - ln_ratio)

    return td
end

"""
    relative_to_specific_humidity(rh::T, t::T, p::T) -> T

Convert relative humidity to specific humidity.
"""
function relative_to_specific_humidity(rh::T, t::T, p::T) where T<:AbstractFloat
    qs = saturation_specific_humidity(t, p)
    q = (rh / T(100)) * qs
    return q
end

"""
    dewpoint_to_specific(td::T, p::T) -> T

Convert dewpoint temperature to specific humidity.
"""
function dewpoint_to_specific(td::T, p::T) where T<:AbstractFloat
    # Saturation vapor pressure at dewpoint
    e = saturation_vapor_pressure(td)

    # Mixing ratio
    w = epsilon * e / (p - e)

    # Specific humidity
    q = w / (T(1) + w)

    return q
end

# Surface Pressure Operators

"""
    surface_pressure_forward(ps_field::AbstractMatrix{T},
                            lon_grid::AbstractVector{T}, lat_grid::AbstractVector{T},
                            lon_obs::T, lat_obs::T;
                            height_correction::Bool=false,
                            z_obs::T=T(0), z_field::Union{AbstractMatrix{T},Nothing}=nothing) -> T

Forward operator for surface pressure observations.

# Arguments
- `ps_field`: 2D surface pressure field [Pa]
- Grid and observation location parameters
- `height_correction`: Apply barometric height correction
- `z_obs`: Observation height [m]
- `z_field`: Model surface height field [m]

# Returns
- Surface pressure at observation location
"""
function surface_pressure_forward(ps_field::AbstractMatrix{T},
                                 lon_grid::AbstractVector{T},
                                 lat_grid::AbstractVector{T},
                                 lon_obs::T, lat_obs::T;
                                 height_correction::Bool=false,
                                 z_obs::T=T(0),
                                 z_field::Union{AbstractMatrix{T},Nothing}=nothing) where T<:AbstractFloat

    # Interpolate surface pressure
    ps_interp = bilinear_interpolate(ps_field, lon_grid, lat_grid, lon_obs, lat_obs)

    # Apply height correction if requested
    if height_correction && !isnothing(z_field)
        z_model = bilinear_interpolate(z_field, lon_grid, lat_grid, lon_obs, lat_obs)
        dz = z_obs - z_model

        # Barometric formula: p(z) = p0 * exp(-g*dz / (R*T))
        # Using standard temperature
        T_std = T(288.15)
        correction = exp(-g * dz / (Rd * T_std))
        ps_interp *= correction
    end

    return ps_interp
end

"""
    surface_pressure_adjoint!(ps_gradient::AbstractMatrix{T},
                             residual::T,
                             lon_grid::AbstractVector{T}, lat_grid::AbstractVector{T},
                             lon_obs::T, lat_obs::T;
                             height_correction::Bool=false,
                             z_obs::T=T(0), z_field::Union{AbstractMatrix{T},Nothing}=nothing)

Adjoint operator for surface pressure observations.
"""
function surface_pressure_adjoint!(ps_gradient::AbstractMatrix{T},
                                  residual::T,
                                  lon_grid::AbstractVector{T},
                                  lat_grid::AbstractVector{T},
                                  lon_obs::T, lat_obs::T;
                                  height_correction::Bool=false,
                                  z_obs::T=T(0),
                                  z_field::Union{AbstractMatrix{T},Nothing}=nothing) where T<:AbstractFloat

    # Adjoint of height correction
    adjusted_residual = residual

    if height_correction && !isnothing(z_field)
        z_model = bilinear_interpolate(z_field, lon_grid, lat_grid, lon_obs, lat_obs)
        dz = z_obs - z_model
        T_std = T(288.15)
        correction = exp(-g * dz / (Rd * T_std))
        adjusted_residual = residual * correction
    end

    # Adjoint of bilinear interpolation
    bilinear_interpolate_adjoint!(ps_gradient, adjusted_residual,
                                  lon_grid, lat_grid, lon_obs, lat_obs)

    return nothing
end

end # module ConventionalObs
